import tushare as ts
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from backtesting import Backtest, Strategy
from backtesting.lib import crossover

# ========== 初始化设置 ==========
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')  
pro = ts.pro_api()

# ========== 数据获取函数 ==========
def get_stock_data(ts_code, start_date, end_date):
    """获取并预处理股票数据"""
    df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
    df = df.sort_values('trade_date')
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    
    # 标准化列名（首字母大写）
    df.rename(columns={
        'open': 'Open',
        'high': 'High',
        'low': 'Low',
        'close': 'Close',
        'vol': 'Volume'
    }, inplace=True)
    
    return df[['Open', 'High', 'Low', 'Close', 'Volume']]

# ========== 获取多支股票数据 ==========
symbols = ['600519.SH', '000001.SZ', '000651.SZ']
stock_data = {symbol: get_stock_data(symbol, '20180101', '20230101') for symbol in symbols}

# ========== 传统量化策略（双均线策略） ==========
class TraditionalStrategy(Strategy):
    n1 = 5    # 短期均线周期
    n2 = 20   # 长期均线周期
    
    def init(self):
        self.ma_short = self.I(lambda x: pd.Series(x).rolling(self.n1).mean(), self.data.Close)
        self.ma_long = self.I(lambda x: pd.Series(x).rolling(self.n2).mean(), self.data.Close)
        self.stop_loss = 0.95  # 5%止损
        self.take_profit = 1.10  # 10%止盈
    
    def next(self):
        current_price = self.data.Close[-1]
        
        if not self.position:
            if crossover(self.ma_short, self.ma_long):
                self.buy(sl=current_price*self.stop_loss, 
                        tp=current_price*self.take_profit)
        else:
            if crossover(self.ma_long, self.ma_short):
                self.position.close()

# ========== 机器学习策略 ==========
def prepare_features(df):
    """特征工程函数"""
    data = df.copy()
    
    # 计算收益率
    data['returns'] = data['Close'].pct_change()
    
    # 技术指标计算
    data['MA5'] = data['Close'].rolling(5).mean()
    data['MA20'] = data['Close'].rolling(20).mean()
    
    # RSI计算
    delta = data['Close'].diff(1)
    gain = delta.clip(lower=0)
    loss = -delta.clip(upper=0)
    avg_gain = gain.rolling(14).mean()
    avg_loss = loss.rolling(14).mean()
    rs = avg_gain / avg_loss
    data['RSI'] = 100 - (100 / (1 + rs))
    
    # MACD指标
    data['EMA12'] = data['Close'].ewm(span=12).mean()
    data['EMA26'] = data['Close'].ewm(span=26).mean()
    data['MACD'] = data['EMA12'] - data['EMA26']
    
    # 波动率
    data['Volatility'] = data['Close'].rolling(20).std()
    
    # 目标变量（次日是否上涨）
    data['target'] = (data['Close'].shift(-1) > data['Close']).astype(int)
    
    # 处理异常值
    data.replace([np.inf, -np.inf], np.nan, inplace=True)
    data = data.ffill()  # 替换fillna(method='ffill')
    
    return data.dropna()

class MLStrategy(Strategy):
    def init(self):
        # 准备特征数据
        df = pd.DataFrame({
            'Open': self.data.Open,
            'High': self.data.High,
            'Low': self.data.Low,
            'Close': self.data.Close,
            'Volume': self.data.Volume
        })
        prepared_data = prepare_features(df)
        
        # 划分训练集和测试集
        split = int(len(prepared_data)*0.8)
        features = ['MA5', 'MA20', 'RSI', 'MACD', 'Volatility']
        X_train = prepared_data[features][:split]
        y_train = prepared_data['target'][:split]
        
        # 训练随机森林模型
        self.model = RandomForestClassifier(n_estimators=100, random_state=42)
        self.model.fit(X_train, y_train)
        
        # 生成预测概率
        self.probabilities = self.model.predict_proba(prepared_data[features])[:, 1]
        self.current_idx = 0
    
    def next(self):
        if self.current_idx >= len(self.probabilities):
            return
        
        current_price = self.data.Close[-1]
        prob = self.probabilities[self.current_idx]
        self.current_idx += 1
        
        position_size = self.position.size if self.position else 0
        buy_threshold = 0.65 - position_size*0.05  # 仓位越大阈值越高
        sell_threshold = 0.35 + position_size*0.05
        
        if not self.position:
            if prob > buy_threshold:
                self.buy(sl=current_price*0.95, 
                        tp=current_price*1.10)
        else:
            if prob < sell_threshold:
                self.position.close()

# ========== 回测系统 ==========
def run_backtest(data, strategy_class):
    """执行回测并返回结果"""
    bt = Backtest(
        data=data,
        strategy=strategy_class,
        cash=1000000,
        commission=0.001,
        exclusive_orders=True
    )
    stats = bt.run()
    return stats

# ========== 执行多股票回测 ==========
results = []
for symbol in symbols:
    print(f"\n正在回测 {symbol}...")
    data = stock_data[symbol]
    
    # 传统策略回测
    traditional_stats = run_backtest(data, TraditionalStrategy)
    
    # 机器学习策略回测
    ml_stats = run_backtest(data, MLStrategy)
    
    results.append({
        'Symbol': symbol,
        'Traditional Return (%)': round(traditional_stats['Return [%]'], 2),
        'ML Return (%)': round(ml_stats['Return [%]'], 2),
        'Traditional Sharpe': round(traditional_stats['Sharpe Ratio'], 2),
        'ML Sharpe': round(ml_stats['Sharpe Ratio'], 2)
    })

# ========== 结果展示 ==========
results_df = pd.DataFrame(results)
print("\n策略表现对比：")
print(results_df.to_string(index=False))

print("\n平均表现：")
print(results_df[['Traditional Return (%)', 'ML Return (%)']].mean().to_string())
print(results_df[['Traditional Sharpe', 'ML Sharpe']].mean().to_string())

bt.plot(resample=False, plot_volume=False, plot_pl=False)
plt.figure(figsize=(10,4))
plt.plot(stats['Drawdown [%]'])
plt.title('Max Drawdown Analysis')

# ========== 绘制回测结果 ==========
if len(symbols) > 0:
    bt = Backtest(stock_data[symbols[-1]], TraditionalStrategy, commission=0.001)
    stats = bt.run()
    bt.plot()
if len(symbols) > 0:
    bt = Backtest(stock_data[symbols[0]], TraditionalStrategy, commission=0.001)
    stats = bt.run()
    bt.plot()
if len(symbols) > 0:
    bt = Backtest(stock_data[symbols[1]], TraditionalStrategy, commission=0.001)
    stats = bt.run()
    bt.plot()